#include <iostream>
#include <iomanip>
#include <vector>
#include <gmssl/pbkdf2.h>
#include <gmssl/rand.h>
#include <gmssl/sm9.h>
#include <gmssl/sm4.h>
using namespace std;
vector<uint8_t> randomData(size_t len){
    vector<uint8_t> res(len);
    rand_bytes(res.data(),len);
    return res;
}
vector<uint8_t> sm2_encrypt2(const SM2_KEY* key, const vector<uint8_t>& in) {
    size_t len=in.size()+SM2_MAX_CIPHERTEXT_SIZE;
    vector<uint8_t> res(len);
    sm2_encrypt(key,in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm2_decrypt2(const SM2_KEY* key,const vector<uint8_t>& in) {
    size_t len=in.size();
    vector<uint8_t> res(len);
    sm2_decrypt(key, in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm4_encrypt2(
    const vector<uint8_t>& key,const vector<uint8_t>& iv,
    const vector<uint8_t>& aad,const vector<uint8_t>& plaintext,
    vector<uint8_t>& tag) {
    SM4_KEY sm4_key;
    sm4_set_encrypt_key(&sm4_key, key.data());
    vector<uint8_t> ciphertext(plaintext.size());
    tag.resize(SM4_GCM_MAX_TAG_SIZE);
    sm4_gcm_encrypt(
        &sm4_key,iv.data(),iv.size(),
        aad.data(),aad.size(),
        plaintext.data(),plaintext.size(),
        ciphertext.data(),
        tag.size(),tag.data()
    );
    return ciphertext;
}
vector<uint8_t> sm4_decrypt2(
    const vector<uint8_t>& key,const vector<uint8_t>& iv,
    const vector<uint8_t>& aad,const vector<uint8_t>& ciphertext,
    const vector<uint8_t>& tag) {
    SM4_KEY sm4_key;
    sm4_set_encrypt_key(&sm4_key, key.data());
    vector<uint8_t> plaintext(ciphertext.size());
    if(sm4_gcm_decrypt(
        &sm4_key,iv.data(),iv.size(),
        aad.data(),aad.size(),
        ciphertext.data(),ciphertext.size(),
        tag.data(),tag.size(),
        plaintext.data())!=1){
        throw runtime_error("GCM标签验证失败");
    }
    return plaintext;
}
SM9_SIGNATURE sm9_sign(const SM9_SIGN_KEY* key,string s){
    SM3_CTX ctx;
    sm3_init(&ctx);
    sm3_update(&ctx,(uint8_t*)s.c_str(),s.size());
    SM9_SIGNATURE sig;
    sm9_do_sign(key,&ctx,&sig);
    return sig;
}
bool sm9_verify(const SM9_SIGN_MASTER_KEY* key,SM9_SIGNATURE sig,string id,string s){
    SM3_CTX ctx;
    sm3_init(&ctx);
    sm3_update(&ctx,(uint8_t*)s.c_str(),64);
    return sm9_do_verify(key,id.c_str(),id.size(),&ctx,&sig)==1;
}
void exchangeSm2Key(const SM2_KEY* a,const SM2_KEY* b,uint8_t *shared){
    SM2_POINT b_pub=b->public_key;
    SM2_POINT S_a;
    sm2_point_mul(&S_a,a->private_key,&b_pub);
    memcpy(shared,S_a.x,32);
    memcpy(shared+32,S_a.y,32);
}
const string BASE64_CHARS="ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
std::string base64_encode(const std::string &input) {
    std::string encoded;
    encoded.reserve((input.size()+2)/3);
    int val=0,valb=-6;
    for(unsigned char c:input) {
        val=(val<<8)+c;
        valb+=8;
        while(valb>=0) {
            encoded.push_back(BASE64_CHARS[(val>>valb)&0x3F]);
            valb-=6;
        }
    }
    if(valb>-6) encoded.push_back(BASE64_CHARS[((val<<8)>>(valb+8))&0x3F]);
    while(encoded.size()%4) encoded.push_back('=');
    return encoded;
}
void printSm2Key(const SM2_KEY& key,string lab){
    cout<<lab+" public key:\n  x:"
        <<base64_encode(string(key.public_key.x,key.public_key.x+32))<<"\n  y:"
        <<base64_encode(string(key.public_key.y,key.public_key.y+32))<<'\n'
        <<lab+" private key:"<<base64_encode(string(key.private_key,key.private_key+32))
        <<"\n\n";
}
int main() {
    // 生成sm9签名密钥
    string id1="Alice7891@sb.com";
    SM9_SIGN_MASTER_KEY amsk;
    SM9_SIGN_KEY akey;
    sm9_sign_master_key_generate(&amsk);
    sm9_sign_master_key_extract_key(&amsk,id1.data(),id1.size(),&akey);
    string id2="Bob114514@sb.com";
    SM9_SIGN_MASTER_KEY bmsk;
    SM9_SIGN_KEY bkey;
    sm9_sign_master_key_generate(&bmsk);
    sm9_sign_master_key_extract_key(&bmsk,id2.data(),id2.size(),&bkey);
    // 生成sm2密钥
    SM2_KEY a_key,b_key;
    sm2_key_generate(&a_key);
    sm2_key_generate(&b_key);
    printSm2Key(a_key,"Alice");
    printSm2Key(b_key,"Bob");
    // Alice(a)和Bob(b)创建签名
    uint8_t a_pub_s[64];
    memcpy(a_pub_s,a_key.public_key.x,32);
    memcpy(a_pub_s+32,a_key.public_key.y,32);
    SM9_SIGNATURE siga=sm9_sign(&akey,string(a_pub_s,a_pub_s+64));
    uint8_t b_pub_s[64];
    memcpy(b_pub_s,b_key.public_key.x,32);
    memcpy(b_pub_s+32,b_key.public_key.y,32);
    SM9_SIGNATURE sigb=sm9_sign(&bkey,string(b_pub_s,b_pub_s+64));
    // Bob(b)和Alice(a)验签
    if(!sm9_verify(&amsk,siga,id1,string(a_pub_s,a_pub_s+64))){
        cout<<"Signature a failed!"<<endl;
        return -1;
    }
    if(!sm9_verify(&bmsk,sigb,id2,string(b_pub_s,b_pub_s+64))){
        cout<<"Signature b failed!\n";
        return -1;
    }
    // 标准ECDH密钥交换
    uint8_t shared_a[64];
    uint8_t shared_b[64];
    // Alice(a)计算共享密钥: S_a = a_priv * B_pub
    exchangeSm2Key(&a_key,&b_key,shared_a);
    // Bob(b)计算共享密钥: S_b = b_priv * A_pub
    exchangeSm2Key(&b_key,&a_key,shared_b);
    // 验证共享密钥一致性
    if(memcmp(shared_a,shared_b,64)!=0) {
        cout<<"失败！\n";
        return 0;
    }
    // 使用共享密钥的中间部分作为私钥
    SM2_KEY key1,key2;
    sm2_key_set_private_key(&key1,shared_a+16);
    sm2_key_set_private_key(&key2,shared_b+16);
    printSm2Key(key1,"shared");
    // 生成sm4加密密钥
    vector<uint8_t> sm4_key=randomData(16);
    vector<uint8_t> sm4_iv=randomData(12);
    vector<uint8_t> sm4_aad=randomData(12);
    cout<<"Sm4 key:"<<base64_encode(string(sm4_key.begin(),sm4_key.end()))<<'\n';
    cout<<"Sm4 iv:"<<base64_encode(string(sm4_iv.begin(),sm4_iv.end()))<<'\n';
    cout<<"Sm4 aad:"<<base64_encode(string(sm4_aad.begin(),sm4_aad.end()))<<"\n\n";
    // 把sm4密钥、iv和aad合起来加密
    vector<uint8_t> tag;
    vector<uint8_t> sm4_keyz=sm4_key;
    sm4_keyz.reserve(40);
    sm4_keyz.append_range(sm4_iv);
    sm4_keyz.append_range(sm4_aad);
    vector<uint8_t> key_enc=sm2_encrypt2(&key1,sm4_keyz);
    // 加密密钥组并拆分
    vector<uint8_t> key_dec=sm2_decrypt2(&key2,key_enc);
    vector<uint8_t> dec_key(key_dec.begin(),key_dec.begin()+16);
    vector<uint8_t> dec_iv(key_dec.begin()+16,key_dec.begin()+28);
    vector<uint8_t> dec_aad(key_dec.begin()+28,key_dec.end());
    cout<<"Dec key:"<<base64_encode(string(sm4_key.begin(),sm4_key.end()))<<'\n';
    cout<<"Dec iv:"<<base64_encode(string(sm4_iv.begin(),sm4_iv.end()))<<'\n';
    cout<<"Dec aad:"<<base64_encode(string(sm4_aad.begin(),sm4_aad.end()))<<"\n\n";
    // 使用sm4对称加密
    string message="我是a!";
    vector<uint8_t> enc=sm4_encrypt2(sm4_key,sm4_iv,sm4_aad,vector<uint8_t>(message.begin(),message.end()),tag);
    cout<<"Enc:"<<base64_encode(string(enc.begin(),enc.end()))<<'\n';
    vector<uint8_t> dec=sm4_decrypt2(sm4_key,sm4_iv,sm4_aad,enc,tag);
    string dec_s(dec.begin(),dec.end());
    cout<<"Dec:"<<dec_s<<endl;
    cout<<"Message:"<<message<<endl;
    return 0;
}